package Psh::Sql;

use strict;

require Getopt::Std;

#require DBI;

$Psh::Sql::dbh = 0;
$Psh::Sql::type = "";
$Psh::Sql::curdb = "";

@Psh::Sql::cmpl_keywords = qw ( select table from update using limit create set alter);
@Psh::Sql::cmpl_list;

# if this is 0, then the strategy Db uses the full line mode instead of token
# (this affects how > and < and other shell tokens get split
$Psh::Sql::full_line = 1;

$Psh::Prompt::prompt_vars{"Q"} = sub{ if ($Psh::Sql::curdb) { return "SQL:$Psh::Sql::curdb" } else { return "psh"} };


my $db_defaults = $ENV{'HOME'} . "/.my.cnf";
my $already_did_aliases = 0;
my %cmpl;
my $elapsed_time = time;
my @dbi_options;

sub has_DBI {
	require DBI;
	if ($@) {
		Psh::Util::print_error("could not load DBI\n");
		return(0);
	}
	return(1);
}


sub db_open
{
        my $line=shift;
        local @ARGV = @{shift()};

        if( ! has_DBI()) {
                Psh::Util::print_error('no DBI module loaded');
                return (0,undef);
        }
        my %opt;
        Getopt::Std::getopts('h:P:u:p:t:d:',\%opt);

        my $host = $Psh::host;
        my $port = 3306;
        my $user = "anonymous";
        my $password = "";
        my $database = "mysql";
        my $dbtype = "mysql";

	# read the ~/.my.cnf for the default settings
	if (-e $db_defaults) {
		if (open DBDEF, "<$db_defaults" ) {
			while (my $line = <DBDEF>) {
				if ($line =~ /(\S+)\s*=\s*(\S+)/) {
					if (lc $1 eq "port") {
						$port = $2;
					} elsif (lc $1 eq "user") {
						$user = $2;
					} elsif (lc $1 eq "database") {
						$database = $2;
					} elsif (lc $1 eq "password") {
						$password = $2;
					} elsif (lc $1 eq "host") {
						$host = $2;
					}

				}
			}
		} else {
			Psh::Util::print_debug ("Warning: Cannot open db_default: $db_defaults for reading\n");
		}
	}

	# overwrite any default settings with the command line options
        $host = $opt{'h'}     if (defined $opt{'h'});
        $port = $opt{'P'}     if (defined $opt{'P'});
        $user = $opt{'u'}     if (defined $opt{'u'});
        $password = $opt{'p'} if (defined $opt{'p'});
        $database = $opt{'d'} if (defined $opt{'d'});
        $dbtype = $opt{'t'}   if (defined $opt{'t'});

        my $dsn = "DBI:$dbtype:database=$database;host=$host;port=$port";
        Psh::Util::print_debug ("trying dsn: $dsn\n");
	@dbi_options = ( $dsn, $user, $password, { RaiseError => 0, PrintError => 0} );
        my $dbh = connect_dbh(@dbi_options);
	return (0,undef) if (! $dbh or ! $dbh->{Active});

        # save database handle in Psh::Sql namespace
        $Psh::Sql::dbh = $dbh;
        $Psh::Sql::type = $dbtype;
	$Psh::Sql::curdb = $database;

	if (! $already_did_aliases) { 
		setup_aliases( "select create update alter" );
		$already_did_aliases = 1;
	}

        return(1,undef);

}

# gets the dbh and connects or reconnects if missing
sub get_dbh {
        my $dbh = $Psh::Sql::dbh;
        my $curdb = $Psh::Sql::curdb;
        if( ! $dbh) {
		# never opened, try defaults
                Psh::Util::print_debug("Warning: no open database, trying defaults\n");
		unless (db_open()) {
                	Psh::Util::print_error("Error: could not open using defaults\n");
                	return (0,undef);
		}
		$dbh = $Psh::Sql::dbh;
		if (! $dbh) {
                	Psh::Util::print_error("Error: dbh empty, could not open using defaults\n");
                	return (0,undef);
		}
        }

	if (! $dbh->{Active}) {
		# opened but inactive, try reopening
                Psh::Util::print_debug("Warning: database handle inactive, reconnecting\n");
		my $old_dsn = $dbi_options[0];
		my $new_dsn = $old_dsn;
		$new_dsn =~ s/database=\S+?\;/database=$curdb\;/;
		$dbi_options[0] = $new_dsn;
		unless ($dbh = connect_dbh(@dbi_options)) {
                	Psh::Util::print_error("Error: could not open using previous set of options with new db $curdb\n");
			$dbi_options[0] = $old_dsn;
			unless ($dbh = connect_dbh(@dbi_options)) {
                		Psh::Util::print_error("Error: could not open using previous set of options.\n");
                		return (0,undef);
			}
		}
	}

        $Psh::Sql::dbh = $dbh;
	return($dbh);
}

sub connect_dbh {
	my (@connect_options) = @_;
	my $dbh;
        if (! ($dbh = DBI->connect(@connect_options))) {
                Psh::Util::print_error ("Error: Cannot connect to database [" . 
		        $dbh->errstr . 
		        "] using options:\n\t". join("\t\n", @connect_options). "\n");
                return( 0,undef);
        } else {
                Psh::Util::print_debug ("[Connected to database]\n");
        }
	return($dbh);
}

sub prepare_and_execute {
	my ($statement, $silent) = @_;
	my ($dbh, $sth);
	unless ($dbh = get_dbh()) {
		Psh::Util::print_error("Cannot get database handle\n");
		return(0,undef);
	}
        Psh::Util::print_out( "SQL: $statement\n") unless ($silent);
        $sth = $dbh->prepare($statement);
        #eval { $sth->execute()  };
        $sth->execute();
	if ($sth->err) {
		if ($sth->errstr =~ /mysql server has gone away/i) {
			Psh::Util::print_debug("server has gone away, reseting connection\n") unless ($silent);
			$dbh->disconnect();
			unless ($dbh = get_dbh()) {
				Psh::Util::print_error("could not reset connection\n");
				return(0,undef);
			}
			$sth = $dbh->prepare($statement);
			#eval { $sth->execute(); };
			$sth->execute();
			if ($sth->err and $sth->errstr) {
				Psh::Util::print_error("SQLerror on statement with new connection:" .  $sth->errstr() . "\n");
				return(0,undef);
			}
		}
		if ($sth->errstr ) {
			Psh::Util::print_error("SQL ERROR " .  $sth->err() . " : " . $sth->errstr() . "\n");
			return(0,undef);
		}
	}
	return($sth);
}

# setup aliases and new completion functions for those aliases
sub setup_aliases
{
	my ($text) = @_;

	foreach my $dbcom ( split( ' ', $text) ) {
		# set the alias using the builtin alias setup
		my $text = "alias $dbcom='db $dbcom'";
		Psh::evl($text,'built_in','fallback_builtin');

		# set the completion functions using the builtin command: complete
		$text = "complete -F Psh::Builtins::Dbcore::cmpl_dbcore $dbcom";
		Psh::evl($text,'built_in','fallback_builtin');
	}
}

sub pretty_print {
	my ($statement_handle, $truncate_fields) = @_;
        my (@field_widths, @all_rows, $cell_value, @field_names, $field_name, $i);
	my $total_num_rows = 0;

	$truncate_fields = 0 if (! $truncate_fields);

	# get the field names, and initialize the field_widths
	@field_names = @{$statement_handle->{NAME}};
       	my $total_row_width = 0;
	for $i ( 0 .. $#field_names) {
		$field_widths[$i] = length($field_names[$i]);
		$total_row_width += $field_widths[$i];
	}

	# save them as the first row for display purposes
	push(@all_rows, \@field_names);

	# save all the data for the table, and find the widest entry for each field
        while ( my $row_array = $statement_handle->fetchrow_arrayref ) {
		$total_num_rows++;
                my @row_data;
                foreach $i ( 0 .. $#$row_array ) {
                        $cell_value = $row_array->[$i];
			if ($truncate_fields) {
				$cell_value =~ s/\n/\\n/g;
				$cell_value =~ s/\t/\\t/g;
			}
                        push(@row_data, $cell_value);
                        if ( length($cell_value) > $field_widths[$i] ) {
                                $field_widths[$i] = length($cell_value);
                        }
                }
                push(@all_rows, \@row_data);
        }

	# calculate the total row width and truncate any extra wide columns
	if ($truncate_fields) {
		# adjust the max width to include the spaces between columns and the edges
		my $max_width = 100;
		my $borders = ( 2 + (3 * ($#field_widths + 1)));

		# total_row_width is currently storing the width of all the column names as printed
		# so, only try to truncate fields if the column names are narrow enough.
		if ( $total_row_width + $borders < $max_width) {
			$total_row_width = 0;
        		foreach $i ( 0 .. $#field_widths) {
				$total_row_width += $field_widths[$i];
        		}
			if ($total_row_width + $borders > $max_width) {
				foreach $i ( sort {$field_widths[$b] <=> $field_widths[$a] } (0 .. $#field_widths) ) {
        Psh::Util::print_debug( "debug: checking row: $i max_width=$max_width total=$total_row_width cur=$field_widths[$i] borders=$borders\n");
					if ( $field_widths[$i] > 0.25 * $total_row_width) {
						my $width_of_other_cols = $total_row_width - $field_widths[$i];
						$total_row_width -= $field_widths[$i];
						if ($max_width > $width_of_other_cols + $borders) {
							$field_widths[$i] = $max_width - $borders - $width_of_other_cols;
        Psh::Util::print_debug( "\tdebug: $i other=$width_of_other_cols new=$field_widths[$i]\n");
						} else {
							$field_widths[$i] = int(($max_width -$borders)/($#field_widths+1));
        Psh::Util::print_debug( "\tdebug: $i new=$field_widths[$i]\n");
						}
						$total_row_width += $field_widths[$i];
					}
					last if ($total_row_width + $borders < $max_width);
				}
			} else {
				$truncate_fields = 0;
			}
		} else {
			$truncate_fields = 0;
		}
	}

	# this is a row full of dashes
        my $all_dashes =  "+";
        foreach $i ( 0 .. $#field_widths) {
        	$all_dashes .=  "-" x (2 + $field_widths[$i]) . "+";
        }

        # pretty print the table
        Psh::Util::print_out( "$all_dashes\n");
        my $first_row = 1;
        foreach my $row_data_ref (@all_rows) {

                my $row_string = "| ";
                foreach $i ( 0 .. $#$row_data_ref) {
                        $cell_value = $row_data_ref->[$i];
			$cell_value = substr($cell_value, 0, ($field_widths[$i]-3)) . '...' 
					if ($truncate_fields and length($cell_value) > $field_widths[$i]);
                        $row_string .= $cell_value;
                        $row_string .= " " x ($field_widths[$i] - length($cell_value) );
                        $row_string .= " | ";
                }
                Psh::Util::print_out("$row_string\n");
                if ($first_row) {
                        Psh::Util::print_out( "$all_dashes\n");
                        $first_row = 0;
                }
        }

	Psh::Util::print_out( "$all_dashes\n");
	$elapsed_time = time - $elapsed_time;
	$elapsed_time = '<1' if ($elapsed_time < 1);
	$total_num_rows .= " row";
	$total_num_rows .= "s" if ($total_num_rows > 1);
	Psh::Util::print_out( "$total_num_rows in set ($elapsed_time sec)\n");

	return (1,undef);

}


sub db_head
{
        my $line=shift;
        my ($num_rows, $table_name) = @{shift()};

        if (! $num_rows) {
                Psh::Util::print_error("missing table name\n");
                return (0,undef);
        } elsif (! $table_name ) {
                $table_name = $num_rows;
                $num_rows = "-10";
        }
        Psh::Util::print_debug( "debug: line=$line num_rows=$num_rows table=$table_name\n");

        if ($num_rows =~ /^\-(\d*\,{0,1}\d+)/) {
                $num_rows = $1;
        } else {
                Psh::Util::print_error("improperly formatted limit\n");
                return (0,undef);
        }

        if (! $table_name) {
                Psh::Util::print_error("no table entered\n");
                return (0,undef);
        }

	$elapsed_time = time;
	if ($table_name !~ /^\`\S+\`$/) {
		$table_name =~ s/\./\`\.\`/;  # add the backtics on either side of an internal period
		$table_name = '`' . $table_name . '`';  # add the backtics around the table name
	}
	my $sth = prepare_and_execute("select * FROM $table_name LIMIT $num_rows" );
	if (! $sth) {
		return(0,undef);
	} else {
        	Psh::Sql::pretty_print($sth, "truncate_fields");
		$sth->finish();
	}
}

sub cmpl_sql
{
	my( $text, $starttext) = @_;
	my @words= split ' ',$starttext;
	$text =~ s/(.*\,)//; # remove anything in the current word up to a comma
	$starttext .= " $1";
	$Psh::Completion::ac=' ';
	# text has the current word and starttext has the beginning of the line
	#Psh::Util::print_debug("[$text,$starttext]");
	return (1,grep { Psh::Util::starts_with($_,$text) }
		completions($text, $starttext) );
}

sub completions
{
	my ($cur_word, $line) = @_;
	#Psh::Util::print_debug("[$line][$cur_word]");

	if ($cur_word =~ /\./) {
		# if the currently typed word has a period, then only show tables and fields with periods
		return( 
			&get_cmpl_array("dot_fields", $cur_word),
			&get_cmpl_array("dot_tables", $cur_word),
			);
	} elsif ($line =~ /^db\s*use\s*/ ) {
		# starting with db use means we are looking for databases
		return( 
			&get_cmpl_array("databases"),
			);
	} elsif ($line =~ /\swhere\s/ ) {
		# after a "where" means that we are needing fields, not tables
		return( 
			@Psh::Sql::cmpl_keywords,
			&get_cmpl_array("databases"),
			&get_cmpl_array("fields"),
			&get_cmpl_array("tables"),
			);
	} elsif ($line =~ /\sfrom\s$/ ) {
		# after a "from" means that we are only needing tables
		return( 
			@Psh::Sql::cmpl_keywords,
			&get_cmpl_array("databases"),
			&get_cmpl_array("tables"),
			);
	} else {
		# otheriwse return everything that does not have a dot
		return( 
			@Psh::Sql::cmpl_keywords,
			&get_cmpl_array("fields"),
			&get_cmpl_array("tables"),
			);
	}

}

sub db_rehash
{
	my ($database) = @_;
	my ($curdb);
	if (! defined $database or ! $database) {
		$curdb = $Psh::Sql::curdb;
	}
	@{$cmpl{$curdb}{fields}}     = get_fields();
	@{$cmpl{$curdb}{tables}}     = get_tables();
	@{$cmpl{$curdb}{databases}}  = get_databases();
	@{$cmpl{$curdb}{dot_fields}} = get_dot_fields();
	@{$cmpl{$curdb}{dot_tables}} = get_dot_tables();
	return(1,undef);
}

# get the current "type" array or run get_type
# e.g. get @$cmpl{$database}{fields} or run get_fields()
sub get_cmpl_array
{
	my ($type, $word_match) = @_;
	my (@type_array);
	my $curdb = $Psh::Sql::curdb;

	# only send the match args if working on the dot_ subs
	if ($word_match and $type =~ /^dot_/) {
		# just send the word in front of the dot
		if ($type eq "dot_fields" and $word_match =~ /^(\S+?)\.(\S+?)\./) {
			$word_match = " '$1','$2' ";
		} elsif ( $word_match =~ /^(\S+?)\./ ) {
			$word_match = " '','$1' ";
		}
		delete $cmpl{$curdb}{$type};
	} else {
		$word_match = ""; 
	}

	# make the string which is the subroutine call to be eval'ed
	my $sub_name = '&' . "get_$type" . '('. $word_match .')';
	if ( ! defined $cmpl{$curdb}{$type} ) {
		#Psh::Util::print_debug("[running $sub_name]\n");
		@{$cmpl{$curdb}{$type}} = eval ($sub_name);
		#return( eval ($sub_name) ) ;
	} else {
		#Psh::Util::print_debug("[returning $type]\n");
	}
	return( @{$cmpl{$curdb}{$type}} );
}
	
sub get_fields
{
	my ($table_name, $database) = @_;
	my (@field_list, @table_list);

	# if no table is specified it will get all fields from all tables in the current database
	# otherwise it will just return the fields for the table requested
	if ($table_name) {
		@table_list = ($table_name);
	} else {
		@table_list = get_tables();
	}

	if (! defined $database or ! $database) {
		$database = "";
	} else {
		$database = '`' . $database . '`.';
	}

	foreach $table_name ( @table_list ) {
		my $sth = prepare_and_execute("describe $database`$table_name`", "silent");
		while ( my $row_ref = $sth->fetchrow_arrayref ) {
			# in mysql, the first field from a describe statement is the field name
			push(@field_list, $row_ref->[0]);
			
		}
		$sth->finish();
	}

	return(@field_list);
}

sub get_tables
{
	my ($database) = @_;
	my (@table_list);

	if (! defined $database or ! $database) {
		$database = "";
	} else {
		$database = ' from `' . $database. '`';
	}

	my $sth = prepare_and_execute("show tables$database", "silent");
	while ( my $row_ref = $sth->fetchrow_arrayref ) {
		push(@table_list, $row_ref->[0]);
		
	}
	$sth->finish();

	return(@table_list);
}

sub get_databases
{
	my (@db_list);

	my $sth = prepare_and_execute("show databases", "silent");
	while ( my $row_ref = $sth->fetchrow_arrayref ) {
		push(@db_list, $row_ref->[0]);
		
	}
	$sth->finish();

	return(@db_list);
}

sub get_dot_fields
{
	my ($database, $table_match) = @_;
	my (@field_list, $curdb);
	if (! defined $database or ! $database) {
		$curdb = $Psh::Sql::curdb;
	} else {
		$curdb = $database;
	}
	foreach my $table_name ( get_tables($database) ) {
		next if ($table_match and $table_name !~ /$table_match/);
		foreach my $field_name ( get_fields($table_name, $database) ) {
			#push(@field_list, $field_name);
			push(@field_list, $table_name . "." . $field_name);
			push(@field_list, $curdb . "." . $table_name . "." . $field_name);
		}
	}
	return(@field_list);
}

sub get_dot_tables
{
	my ($database, $db_match) = @_;
	my (@db_list, @table_list);
	if (! defined $database or ! $database) {
		@db_list = get_databases();
	} else {
		push(@db_list, $database);
	}
	foreach my $curdb ( @db_list ) {
		next if ($db_match and $curdb !~ /$db_match/);
		foreach my $table_name ( get_tables($curdb) ) {
			#push(@table_list, $table_name);
			push(@table_list, $curdb . "." . $table_name);
		}
	}
	return(@table_list);
}

sub db_use
{
        my $line=shift;
        local @ARGV = @{shift()};

	my $database = $ARGV[0];

	my $sth = prepare_and_execute("use $database") or return(0,undef);

	$sth->finish();
	$Psh::Sql::curdb = $database;
	db_rehash($database);
	return(1,undef);
}

sub db_ls
{
        my $line=shift;
        local @ARGV = @{shift()};

        my (%opt, $opt_long, $opt_rev, $opt_time, $path, $opt_wide);
        Getopt::Std::getopts('lrwtA',\%opt);

	$opt_wide = $opt_rev = $opt_time = 0;
	$opt_long = $opt_time = 1;
	$opt_long = 1 if ($opt{'l'});
	$opt_rev  = 1 if ($opt{'r'});
	$opt_time = 0 if ($opt{'A'});
	$opt_wide = 1 if ($opt{'w'});
	$opt_long = 0 if ($opt_wide);

	#Psh::Util::print_debug("opts: $opt_long $opt_rev $opt_time , remain: " . join(" ",@ARGV) . "\n");

	# add the current database if none are listed
	if ($#ARGV < 0) {
		push (@ARGV, ".");
	}

	foreach $path (@ARGV) {

		my ($statement, @fields, $time_field, $path_prefix, $name_field);
		$path_prefix = $path;
		$path_prefix = "" if ($path =~ /\*/ or $path eq ".");

		if (lc $Psh::Sql::type eq "mysql") {
			# this needs to be reset each time, in case the "/" entry overwrites it
			@fields = qw (Data_length Rows Update_time Name);
			$time_field = "Update_time";
			$name_field = "Name";
			if ($path eq ".") {
				# current database
				$statement = "show table status";
			} elsif ($path eq "/") {
				# show all databases
				$statement = "show databases";
				@fields = qw (Database);
				$time_field = "";
				$name_field = "Database";
			} elsif ($path =~ m|^/(\S+)[/\.](\S+)| ) {
				# show a particular table in a different database
				$statement = "show table status from $1";
				$path = $2;
				$path_prefix = "/$1/";
			} elsif ($path =~ m|^/(\S+)| ) {
				# show the tables in a different database
				$statement = "show table status from $1";
				$path = ".";
			} else {
				# show a particular table in the current database
				$statement = "show table status";
			}
		} else {
        		Psh::Util::print_error( "ls not supported with this db type\n");
			return(0,undef);
		}
			
		my $sth = prepare_and_execute($statement) or return(0,undef);

		# convert the requested path into a perl regex
		my $regex_path = $path;
		$regex_path =~ s/\*/\.\*/g;
		$regex_path =~ s/\?/\./g;

		my (@table_list, %table_info, %max_widths);
		while ( my $row_hash = $sth->fetchrow_hashref ) {
			my ($table_name, $table_time);
			$table_name = $row_hash->{$name_field};
			# 2005-05-09 10:57:08
			$table_time = $row_hash->{$time_field};
			$table_time =~ s/[\s\-\:]//g;
			foreach my $field_name ( @fields ) {
				my $curval = $row_hash->{$field_name};
				$table_info{$table_name}{"fields"}{$field_name} = $curval;
				if (! defined $max_widths{$field_name} or
					$max_widths{$field_name} < length($curval)) {
					$max_widths{$field_name} = length($curval);
				}
			}

			# check if this table matches the requested path (for * and ? in query)
			if ($path eq "." or $path eq "*" or $path eq "/" or
			    lc $table_name eq lc $path or
			    $table_name =~ /^$regex_path$/i ) {

				# save the name, time and other info for each requested table
				$table_info{$table_name}{"time"} = $table_time;
				push(@table_list, $table_name);
			} else {
				# otherwise, delete any previously saved info for this table
				delete $table_info{$table_name};
			}
		}
		$sth->finish();

		# sort the tables
		@table_list = sort { $table_info{$a}{"time"} <=> $table_info{$b}{"time"}  } @table_list
			if ($opt_time);
		@table_list = reverse @table_list
			if ($opt_rev);

		# print the tables
		if ($opt_long) {
			foreach my $table_name ( @table_list ) {
				Psh::Util::print_out("   ");
				foreach my $i ( 0 .. $#fields ) {
					my $field_name = $fields[$i];
					my $curval = $table_info{$table_name}{"fields"}{$field_name};
					Psh::Util::print_out("  ");
					Psh::Util::print_out(" " x ($max_widths{$field_name} - length($curval))) 
						unless ($i == $#fields);
					Psh::Util::print_out($curval);
				}
				Psh::Util::print_out("\n");
			}
		} else {
			Psh::Util::print_list(@table_list);
		}
	}

	return(1,undef);
}

sub db_run
{
        my $line=shift;

        my $dbh = $Psh::Sql::dbh;

        if( ! $dbh) {
                Psh::Util::print_error("Error: you need to run dbopen first\n");
                return (0,undef);
        }

	if ($line =~ /select/) {
		$line =~ s/ \.gt / \> /g;
		$line =~ s/ \.lt / \< /g;
		$line =~ s/ \.ge / \>= /g;
		$line =~ s/ \.le / \<= /g;
	}

	$elapsed_time = time;
	my $sth = prepare_and_execute($line) or return(0,undef);
	if ($line =~ /^select/i or $line =~ /^show/i or $line =~ /^desc/i ) {
        	Psh::Sql::pretty_print($sth);
	}
	# } else {
        #	my $sth = $dbh->do($line);
	#	if ($dbh->err) {
	#		Psh::Util::print_error("SQL Error:". $dbh->errstr() . "\n" );
	#		return (0,undef);
	#	}
	# }
	$sth->finish();
		
	return (1, undef);
}

1;

__END__

